package com.glodon.paas.job.agent.mq;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.atomic.AtomicInteger;

import javax.jms.Destination;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.MessageConsumer;
import javax.jms.Session;

import org.springframework.jms.listener.DefaultMessageListenerContainer;

import com.glodon.paas.job.JobConstants;
import com.glodon.paas.job.agent.worker.WorkerConfig;

public class MultipleDestinationMessageListenerContainer extends DefaultMessageListenerContainer {
    private List<MultipleDestinationMessageConsumer> consumers = new ArrayList<MultipleDestinationMessageConsumer>();
    private Map<String, WorkerConfig> typeWorkerConfigs = new HashMap<String, WorkerConfig>();
    private Map<String, AtomicInteger> concurrent = new HashMap<String, AtomicInteger>();
    private long receiveTimeout = DEFAULT_RECEIVE_TIMEOUT;

    public MultipleDestinationMessageListenerContainer() {
    }

    protected javax.jms.MessageConsumer createConsumer(Session session, Destination destination) throws JMSException {
        if (destination instanceof MultipleDestination) {
            MultipleDestination dest = (MultipleDestination) destination;
            WorkerConfig[] workerConfigs = dest.getWorkerConfigs();
            Map<String, MessageConsumer> consumerMap = new HashMap<String, MessageConsumer>();
            for (WorkerConfig workerConfig : workerConfigs) {
                typeWorkerConfigs.put(workerConfig.getType(), workerConfig);
                concurrent.put(workerConfig.getType(), new AtomicInteger());
                MessageConsumer consumer = session.createConsumer(session.createQueue(workerConfig.getQueueName()));
                consumerMap.put(workerConfig.getType(), consumer);
            }
            MultipleDestinationMessageConsumer consumer = new MultipleDestinationMessageConsumer(consumerMap);
            consumers.add(consumer);
            return consumer;
        } else {
            throw new RuntimeException("destination isn't MultipleDestination");
        }
    }

    public void remove(String type) {
        typeWorkerConfigs.remove(type);
        concurrent.remove(type);
        for (MultipleDestinationMessageConsumer consumer : consumers) {
            consumer.remove(type);
        }
    }

    @Override
    protected Message receiveMessage(MessageConsumer consumer) throws JMSException {
        if (consumer instanceof MultipleDestinationMessageConsumer) {
            MultipleDestinationMessageConsumer multiDestConsumer = (MultipleDestinationMessageConsumer) consumer;
            List<String> ignoreTypes = new ArrayList<String>();
            for (Entry<String, WorkerConfig> entry : typeWorkerConfigs.entrySet()) {
                String type = entry.getKey();
                WorkerConfig workerConfig = entry.getValue();
                int count = concurrent.get(type).get();
                if (count >= workerConfig.getConcurrent()) {
                    ignoreTypes.add(type);
                }
            }
            return (this.receiveTimeout < 0 ? multiDestConsumer.receive(ignoreTypes) : multiDestConsumer.receive(
                    this.receiveTimeout, ignoreTypes));
        } else {
            return super.receiveMessage(consumer);
        }
    }

    @Override
    protected void invokeListener(Session session, Message message) throws JMSException {
        AtomicInteger count = null;
        try {
            String type = message.getStringProperty(JobConstants.PROPERTY_TASK_TYPE);
            count = concurrent.get(type);
            count.incrementAndGet();
            super.invokeListener(session, message);
        } finally {
            if (count != null)
                count.decrementAndGet();
        }
    }

    @Override
    public void setReceiveTimeout(long receiveTimeout) {
        this.receiveTimeout = receiveTimeout;
        super.setReceiveTimeout(receiveTimeout);
    }
}
